import re
import pandas as pd
from sklearn.model_selection import train_test_split
import torch
from torch.optim import AdamW
from torch.utils.data import Dataset, DataLoader
from transformers import BertTokenizer, BertForSequenceClassification
from tqdm import tqdm

# 重新读取数据并设置列名
data = pd.read_csv('data/data.csv', header=None)
data.columns = ['text', 'label', 'Unnamed: 2']

# 对标签进行编码
label_map = {'技术问题': 0, '售前咨询': 1, '投诉': 2, '聊天': 3}

# 使用正则表达式清理标签
def clean_label(label):
    return re.sub(r'[^\w]+', '', label)

data['label'] = data['label'].apply(lambda x: label_map[clean_label(x)])

# 划分数据集
train_texts, test_texts, train_labels, test_labels = train_test_split(data['text'], data['label'], test_size=0.2, random_state=42)
train_texts, val_texts, train_labels, val_labels = train_test_split(train_texts, train_labels, test_size=0.2, random_state=42)

# 加载预训练的 BERT 分词器
tokenizer = BertTokenizer.from_pretrained('model/bert-base-chinese/bert-base-chinese')

# 创建自定义数据集类
class CustomDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_length):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = str(self.texts.iloc[idx])
        label = self.labels.iloc[idx]

        encoding = self.tokenizer.encode_plus(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            truncation=True,
            padding='max_length',
            return_attention_mask=True,
            return_tensors='pt',
        )

        return {
            'input_ids': encoding['input_ids'].flatten(),
            'attention_mask': encoding['attention_mask'].flatten(),
            'labels': torch.tensor(label, dtype=torch.long)
        }

# 定义数据集和数据加载器
max_length = 128
train_dataset = CustomDataset(train_texts, train_labels, tokenizer, max_length)
val_dataset = CustomDataset(val_texts, val_labels, tokenizer, max_length)
test_dataset = CustomDataset(test_texts, test_labels, tokenizer, max_length)

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)

# 加载预训练的 BERT 模型并进行微调
model = BertForSequenceClassification.from_pretrained('model/bert-base-chinese/bert-base-chinese', num_labels=len(label_map))
optimizer = AdamW(model.parameters(), lr=2e-5)

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# 训练模型
Epochs = 100
for epoch in range(Epochs):
    model.train()
    total_loss = 0
    progress_bar = tqdm(train_loader, desc=f'Epoch {epoch + 1}', leave=False)
    for batch in progress_bar:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        model.zero_grad()
        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
        loss = outputs.loss
        total_loss += loss.item()

        loss.backward()
        optimizer.step()

        progress_bar.set_postfix({'loss': total_loss / (progress_bar.n + 1)})

    # 验证模型
    model.eval()
    val_loss = 0
    val_correct = 0
    val_total = 0
    with torch.no_grad():
        for batch in val_loader:
            input_ids = batch['input_ids'].to(device)
            attention_mask = batch['attention_mask'].to(device)
            labels = batch['labels'].to(device)

            outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
            loss = outputs.loss
            val_loss += loss.item()

            _, preds = torch.max(outputs.logits, dim=1)
            val_total += labels.size(0)
            val_correct += (preds == labels).sum().item()

    val_accuracy = val_correct / val_total
    print(f'Epoch {epoch + 1}: Train Loss: {total_loss / len(train_loader)}, Val Loss: {val_loss / len(val_loader)}, Val Accuracy: {val_accuracy}')

# 测试模型
model.eval()
test_correct = 0
test_total = 0
with torch.no_grad():
    for batch in test_loader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels = batch['labels'].to(device)

        outputs = model(input_ids, attention_mask=attention_mask, labels=labels)

        _, preds = torch.max(outputs.logits, dim=1)
        test_total += labels.size(0)
        test_correct += (preds == labels).sum().item()

test_accuracy = test_correct / test_total
print(f'Test Accuracy: {test_accuracy}')

# 保存模型
model.save_pretrained('weights/sort_model')
tokenizer.save_pretrained('weights/sort_model')